# Copyright (c) 2017 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import gc
import logging
import json
from enum import Enum
from typing import List, Optional, Union
from collections import deque

from tornado import gen, httpclient
from tornado.websocket import websocket_connect

import coin.flow.archive as archive
from coin.base.datetime_util import to_datetime
from coin.base.param_util import to_list
from coin.base.timestamp import get_timestamp
from coin.flow.topic_map import get_default_topic_map


class TopicEventPublisher(object):
  def __init__(self, queue, event_publish_topic: Union[str, int], topic_map=None):
    topic_map = topic_map or get_default_topic_map()
    try:
      topic_id = topic_map.search_topic_id(event_publish_topic)[0][1]
    except IndexError:
      raise ValueError('no topic string: %s' % event_publish_topic)
    self._writer = queue.get_writer(topic_id)
    self._logger = logging.getLogger(__name__)

  def __del__(self):
    self.close()

  def close(self):
    if self._writer:
      self._writer.close()
      self._writer = None

  def publish(self, event_type: str, target_topic: Union[str, List[str]], timestamp=None, **kwargs):
    if not self._writer:
      self._logger.exception('No writer')
      return
    timestamp = archive.to_timestamp_int(timestamp or get_timestamp())
    target_topic = to_list(target_topic)
    for topic in target_topic:
      msg = {**{'topic': topic, 'type': event_type}, **kwargs}
      # Keeping topic_seq to be zero allows other writers to write with
      # the same topic_id.
      self._writer.reset_topic_seq()
      self._writer.write(timestamp, msg)


class WebsocketJsonSubscriberBase(object):
  BookFeedDisconnectWindowInNs = 2 * 60 * 1e9  # 2 mins

  class State(Enum):
    INIT = 1
    RUNNING = 2
    CLOSED = 3
    STOP = 4

  def __init__(self, close_callback, logger):
    self._logger = logger
    self._close_callback = close_callback
    self._teardown_callback = []

    self._state = self.State.INIT
    self._url = None
    self._ws = None
    self._ioloop = None

    self._topic_event_publisher = None
    self._target_topic = None

  def __del__(self):
    self.stop()

  def _add_teardown_callback(self, callback):
    self._teardown_callback.append(callback)

  def _enable_topic_event_publisher(self, queue, event_publish_topic, target_topic):
    assert self._topic_event_publisher is None
    self._topic_event_publisher = \
        TopicEventPublisher(queue, event_publish_topic)
    self._add_teardown_callback(self._topic_event_publisher.close)

    topic_map = get_default_topic_map()
    for topic_str in target_topic:
      res = topic_map.search_topic_id(topic_str)
      if not res:
        raise ValueError('Cannot find topic id for topic_string=%s' % topic_str)

    self._target_topic = target_topic

  def _publish_topic_event(self, event_type, **kwargs):
    if self._topic_event_publisher is not None:
      self._topic_event_publisher.publish(event_type, self._target_topic, **kwargs)

  def _register_url(self, url, headers=None):
    self._url = httpclient.HTTPRequest(url=url, headers=headers)

  @property
  def stopped(self):
    return self._state == self.State.STOP

  @gen.coroutine
  def on_open(self):
    self._publish_topic_event('START')

  def on_message(self, timestamp: int, message):
    raise NotImplementedError()

  def _notify_close(self):
    if self._state != self.State.RUNNING:
      return

    self._state = self.State.CLOSED
    self._publish_topic_event('STOP', reason='CONNECTION_CLOSED')
    self._logger.error('Connection closed: %s', self.__class__.__name__)
    if self._ws is not None:
      self._ws.close()
      self._ws = None

    if self._close_callback is not None:
      try:
        self._close_callback()
      except Exception:
        self._logger.exception('Close callback raised an exception')

  def _decode(self, ws_msg):
    return json.loads(ws_msg)

  @gen.coroutine
  def _run_loop(self):
    assert self._url
    try:
      if self._state != self.State.RUNNING:
        return
      try:
        self._ws = yield websocket_connect(self._url)
      except httpclient.HTTPError:
        self._logger.exception('Failed to connect %s' % self._url)
        return
      if self._state != self.State.RUNNING:
        return

      self._logger.info('Connected: %s', self.__class__.__name__)
      yield self.on_open()

      while self._state == self.State.RUNNING:
        ws_msg = yield self._ws.read_message()
        if self.stopped:
          self._logger.info('Subscriber (%s, %d) is stopped!', self.__class__.__name__, id(self))
          break
        timestamp = get_timestamp()
        # Bithumb (maybe some others) occasionally sent empty message which
        # doesn't mean the connection is closed. We should explicitly check
        # the value is None rather than using `if ws_msg:`
        if ws_msg is None:
          self._logger.error('%s: Websocket is closed', self.__class__.__name__)
          break
        if len(ws_msg) == 0:
          continue
        try:
          message = self._decode(ws_msg)
        except Exception:
          self._logger.exception('Invalid message. Skipping: "%s"', ws_msg)
          continue

        try:
          self.on_message(timestamp, message)
        except Exception:
          self._logger.exception('Exception raised from on_message!')
    finally:
      self._notify_close()

  # We experienced websocket Sec-WebSocket-Key error when making multiple ws
  # connections to Bithumb, and after the error the underlying socket is still
  # open, and bithumb server side begin to push messages, whilst the our
  # client side is not ready to receive. This is very weird since the key
  # computation algorithm is standard and simple.
  #
  # Although we are not sure what causes the problem, but it seems that if
  # we add longer time delay (a few seconds) between initiation of there ws
  # connections, this error will gone, hence we add this delay_time parameter.
  def start(self, ioloop, delay_time=None):
    assert not self.stopped
    assert gen.is_coroutine_function(self.on_open)
    assert not gen.is_coroutine_function(self.on_message)
    assert gen.is_coroutine_function(self._run_loop)

    self._state = self.State.RUNNING
    self._ioloop = ioloop
    if delay_time is None:
      ioloop.add_callback(self._run_loop)
    else:
      ioloop.call_later(delay_time, self._run_loop)

  def stop(self, trigger_close_callback=False):
    if self._state == self.State.STOP:
      return

    if trigger_close_callback:
      self._close_callback()

    try:
      if self._state == self.State.RUNNING:
        self._publish_topic_event('STOP')
        self._logger.info('Stopped: %s', self.__class__.__name__)
        if self._ws is not None:
          self._ws.close()
          self._ws = None
    finally:
      self._state = self.State.STOP
      self._tear_down()

  def _tear_down(self):
    for callback in self._teardown_callback:
      try:
        callback()
      except Exception:
        self._logger.exception('Teardown callback raised an exception')


def get_next_unit_time(unit_td=None, offset_td=None, current_dt=None):
  # TODO(jaewon): Dumb implementation. Reimplement it!
  unit_td = unit_td or datetime.timedelta(hours=1)
  offset_td = offset_td or datetime.timedelta()
  current_dt = current_dt or datetime.datetime.now()
  dt = to_datetime(current_dt.date()) + offset_td
  while dt <= current_dt:
    dt += unit_td
  return dt


class RotatingSubscriber(object):
  class State(Enum):
    INIT = 1
    SUBMIT = 2
    RUNNING = 3
    STOP = 4

  def __init__(self,
               init_subscriber,
               queue,
               unit_td=None,
               offset_td=None,
               cool_down_td=None,
               gc_collect_on_cool_down=True,
               logger=None,
               spam_cool_down_td=None):
    self._logger = logger or logging.getLogger(__name__)

    self._queue = queue
    self._init_subscriber = init_subscriber
    self._get_next_timeout = lambda: get_next_unit_time(unit_td, offset_td)
    self._cool_down_td = cool_down_td or datetime.timedelta(minutes=1)
    self._gc_collect_on_cool_down = gc_collect_on_cool_down
    assert unit_td > self._cool_down_td

    self._ioloop = None
    self._subscriber = None
    self._state = self.State.INIT

    # Make sure it won't be treated as spam condition.
    init_reconnect_dt = datetime.datetime.now() - datetime.timedelta(minutes=2)
    self._latest_reconnect_times = deque([init_reconnect_dt] * 5, maxlen=5)
    self._spam_cool_down_td = spam_cool_down_td or datetime.timedelta(seconds=30)

  def __del__(self):
    self.stop()

  def _clean_up(self):
    if self._subscriber is not None:
      self._subscriber.stop()
      self._subscriber = None

  def _start_subscriber(self):
    if self._state == self.State.STOP:
      return
    try:
      assert self._subscriber is None
      self._subscriber = self._init_subscriber(self._queue, self._ioloop, self.on_close)
      self._state = self.State.RUNNING
    except Exception:
      self._logger.exception('An exception is raised during initializing subscriber')
      self.on_close()

  def _if_spam(self):
    # Reconnect happens more than 5 times a minute,
    # treat as spam and wait for 30 seconds.
    self._latest_reconnect_times.popleft()
    self._latest_reconnect_times.append(datetime.datetime.now())
    if self._latest_reconnect_times[-1] - self._latest_reconnect_times[0] < \
        datetime.timedelta(minutes=1):
      return True
    else:
      return False

  def _submit(self, later_td=None):
    if self._state != self.State.SUBMIT:
      if self._if_spam():
        self._logger.error('Spam cool down!')
        self._ioloop.add_timeout(self._spam_cool_down_td, self._start_subscriber)
      else:
        self._ioloop.add_timeout(later_td or self._cool_down_td, self._start_subscriber)
      self._state = self.State.SUBMIT
      if self._gc_collect_on_cool_down:
        res = gc.collect()
        self._logger.info('%d objects are collected' % res)

  def on_close(self):
    if self._state == self.State.STOP:
      return
    self._clean_up()
    self._submit()

  def on_timeout(self):
    if self._state == self.State.STOP:
      return
    self._clean_up()
    self._submit()
    self._ioloop.add_timeout(self._get_next_timeout() - datetime.datetime.now(), self.on_timeout)

  def start(self, ioloop, later_td: Optional[datetime.timedelta] = None):
    assert self._state == self.State.INIT
    self._ioloop = ioloop
    if later_td is None or later_td.total_seconds() == 0:
      self._start_subscriber()
    else:
      self._submit(later_td)
    self._ioloop.add_timeout(self._get_next_timeout() - datetime.datetime.now(), self.on_timeout)

  def stop(self):
    self._state = self.State.STOP
    self._clean_up()
